#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import glob
from datetime import datetime
import openpyxl
from openpyxl.utils import get_column_letter
from copy import copy


def get_template_structure(template_file):
    """
    获取模板文件的结构，包括标题行位置和内容
    返回字典：{工作表名: {标题行内容: 标题行索引}}
    """
    print(f"正在读取模板文件: {template_file}")
    template_wb = openpyxl.load_workbook(template_file)
    headers_info = {}
    
    for sheet_name in template_wb.sheetnames:
        sheet = template_wb[sheet_name]
        headers_info[sheet_name] = {}
        
        # 从第4行开始查找标题行（合并单元格的行）
        row_idx = 4
        while row_idx <= sheet.max_row:
            # 检查当前行是否是合并单元格
            is_merged_row = False
            for merged_range in sheet.merged_cells.ranges:
                if merged_range.min_row == row_idx and merged_range.max_row == row_idx:
                    is_merged_row = True
                    # 获取合并单元格的内容
                    header_text = sheet.cell(row=row_idx, column=merged_range.min_col).value
                    if header_text:
                        headers_info[sheet_name][header_text] = row_idx
                    break
            
            row_idx += 1
    
    template_wb.close()
    return headers_info


def extract_data_by_sections(data_file):
    """
    从数据文件中按标题行提取数据
    返回字典：{工作表名: {标题行内容: [数据行列表]}}
    """
    print(f"提取数据文件: {data_file}")
    data_wb = openpyxl.load_workbook(data_file)
    extracted_data = {}
    
    for sheet_name in data_wb.sheetnames:
        sheet = data_wb[sheet_name]
        extracted_data[sheet_name] = {}
        
        # 获取所有标题行的位置和内容
        header_rows = []  # [(row_index, header_text)]
        for merged_range in sheet.merged_cells.ranges:
            if merged_range.min_row >= 4:  # 从第4行开始
                header_text = sheet.cell(row=merged_range.min_row, column=merged_range.min_col).value
                if header_text:
                    header_rows.append((merged_range.min_row, header_text))
        
        # 按照标题行顺序排序
        header_rows.sort()
        
        # 提取每个标题行下面的数据
        for i, (header_row, header_text) in enumerate(header_rows):
            # 确定数据结束行
            if i < len(header_rows) - 1:
                end_row = header_rows[i+1][0]  # 下一个标题行
            else:
                end_row = sheet.max_row + 1  # 最后一行+1（因为range不包含结束值）
            
            # 提取数据行（从标题行的下一行开始）
            section_data = []
            for row_idx in range(header_row + 1, end_row):
                # 检查是否是数据行（第一列有内容且是数字）
                cell_value = sheet.cell(row=row_idx, column=1).value
                if cell_value is not None:
                    try:
                        int(str(cell_value).strip())
                        # 是数字，提取整行数据
                        row_data = []
                        for col in range(1, sheet.max_column + 1):
                            cell = sheet.cell(row=row_idx, column=col)
                            row_data.append(cell.value)
                        section_data.append(row_data)
                    except (ValueError, AttributeError):
                        # 不是数字，跳过该行
                        pass
            
            extracted_data[sheet_name][header_text] = section_data
    
    data_wb.close()
    return extracted_data


def clean_and_renumber_sheet(sheet):
    """
    清理空行并重新编号
    从第4行开始，删除A列为空的行，并对非合并单元格行的A列重新从1开始编号
    """
    # 获取所有合并单元格行号
    merged_rows = set()
    for merged_range in sheet.merged_cells.ranges:
        for row in range(merged_range.min_row, merged_range.max_row + 1):
            merged_rows.add(row)
    
    # 从最后一行开始向前处理，这样删除行不会影响前面行的行号
    row_idx = sheet.max_row
    
    # 先删除空行
    while row_idx >= 4:
        # 如果当前行是合并单元格行，直接跳过
        if row_idx in merged_rows:
            row_idx -= 1
            continue
        
        # 检查A列是否为空
        cell_value = sheet.cell(row=row_idx, column=1).value
        if cell_value is None or str(cell_value).strip() == "":
            # A列为空，删除此行
            sheet.delete_rows(row_idx)
        
        row_idx -= 1
    
    # 再进行重新编号
    current_number = 1
    row_idx = 4
    
    while row_idx <= sheet.max_row:
        # 如果当前行是合并单元格行，重置编号
        if row_idx in merged_rows:
            current_number = 1
        else:
            # 非合并单元格行，设置序号
            if sheet.cell(row=row_idx, column=1).value is not None:
                sheet.cell(row=row_idx, column=1).value = current_number
                current_number += 1
        
        row_idx += 1


def merge_excel_files(template_file, data_files, output_file):
    """
    合并Excel文件的主函数
    """
    print("开始合并数据...")
    
    # 读取模板文件
    template_wb = openpyxl.load_workbook(template_file)
    template_headers = get_template_structure(template_file)
    
    # 创建一个新的工作簿用于输出
    output_wb = openpyxl.Workbook()
    # 删除默认工作表
    output_wb.remove(output_wb.active)
    
    # 复制模板工作表到新工作簿
    template_wb_new = openpyxl.load_workbook(template_file)
    for sheet_name in template_wb_new.sheetnames:
        template_sheet = template_wb_new[sheet_name]
        output_sheet = output_wb.create_sheet(title=sheet_name)
        
        # 复制模板的所有内容和格式
        for row in template_sheet.iter_rows():
            for cell in row:
                new_cell = output_sheet[cell.coordinate]
                new_cell.value = cell.value
                
                # 复制样式
                if cell.has_style:
                    new_cell.font = copy(cell.font)
                    new_cell.border = copy(cell.border)
                    new_cell.fill = copy(cell.fill)
                    new_cell.number_format = cell.number_format
                    new_cell.protection = copy(cell.protection)
                    new_cell.alignment = copy(cell.alignment)
        
        # 复制合并单元格信息
        for merged_range in template_sheet.merged_cells.ranges:
            output_sheet.merge_cells(str(merged_range))
    
    template_wb_new.close()
    
    # 收集所有数据文件的数据
    all_extracted_data = {}
    for sheet_name in output_wb.sheetnames:
        all_extracted_data[sheet_name] = {}
        # 初始化所有标题行的数据列表
        for header in template_headers.get(sheet_name, {}):
            all_extracted_data[sheet_name][header] = []
    
    # 处理所有数据文件
    for data_file in data_files:
        extracted_data = extract_data_by_sections(data_file)
        
        # 合并数据
        for sheet_name in extracted_data:
            if sheet_name in all_extracted_data:
                for header in extracted_data[sheet_name]:
                    if header in all_extracted_data[sheet_name]:
                        all_extracted_data[sheet_name][header].extend(extracted_data[sheet_name][header])
    
    # 将数据插入到输出工作簿中
    for sheet_name in output_wb.sheetnames:
        sheet = output_wb[sheet_name]
        
        # 为每个标题行插入数据
        if sheet_name in template_headers:
            # 按照行号升序排列，逐个处理标题行
            sorted_headers = sorted(template_headers[sheet_name].items(), key=lambda x: x[1])
            
            # 记录每个标题行插入数据后新增的行数，用于调整后续标题行的位置
            offset = 0
            
            for header, original_header_row in sorted_headers:
                header_row = original_header_row + offset  # 调整后的标题行位置
                
                if header in all_extracted_data[sheet_name]:
                    # 获取该标题行下的所有数据
                    section_data = all_extracted_data[sheet_name][header]
                    if not section_data:
                        continue
                    
                    print(f"  在 '{header}' 下插入 {len(section_data)} 行数据")
                    
                    # 在标题行下方插入数据行
                    insert_pos = header_row + 1
                    for i, data_row in enumerate(section_data):
                        sheet.insert_rows(insert_pos + i)
                        
                        # 填充数据，保持原始数据的序号列（第1列）
                        for col_idx, value in enumerate(data_row, 1):
                            sheet.cell(row=insert_pos + i, column=col_idx).value = value
                    
                    # 更新偏移量
                    offset += len(section_data)
    
    # 对每个工作表清理空行并重新编号
    for sheet_name in output_wb.sheetnames:
        print(f"清理和重新编号工作表: {sheet_name}")
        clean_and_renumber_sheet(output_wb[sheet_name])
    
    # 保存合并后的文件
    print(f"保存合并文件: {output_file}")
    output_wb.save(output_file)
    output_wb.close()
    template_wb.close()
    
    print("合并完成！")


def main():
    if len(sys.argv) != 2:
        print("使用方法: python merge.py [文件夹路径]")
        print("示例: python merge.py ./0930")
        sys.exit(1)
    
    folder_path = sys.argv[1]
    
    # 检查文件夹是否存在
    if not os.path.exists(folder_path):
        print(f"错误: 文件夹 '{folder_path}' 不存在")
        sys.exit(1)
    
    # 检查模板文件是否存在
    template_file = "moban.xlsx"
    if not os.path.exists(template_file):
        print(f"错误: 模板文件 '{template_file}' 不存在")
        sys.exit(1)
    
    # 获取所有xlsx文件
    data_files = glob.glob(os.path.join(folder_path, "*.xlsx"))
    if not data_files:
        print(f"错误: 在文件夹 '{folder_path}' 中未找到任何xlsx文件")
        sys.exit(1)
    
    print(f"找到 {len(data_files)} 个数据文件")
    
    # 生成输出文件名（当前日期）
    current_date = datetime.now().strftime("%Y%m%d")
    output_file = f"{current_date}.xlsx"
    
    # 执行合并
    merge_excel_files(template_file, data_files, output_file)


if __name__ == "__main__":
    main()